(*

   Copyright 2025, HE Weikun

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

*)

exception Counter_Example of int list

(* `filter_count f l` gives the number of elements in the list l satisfying the predicate f and the the list of elements that does not *)
let rec filter_count f = function
    | [] -> (0, [])
    | a::l ->
      let (n, r) = filter_count f l in
      if f a then (n + 1, r) else (n, a::r)

(* This is the recursive function that enumerates all subset that is the intersection of an upper set with a lower set.
  poset : bool array array
  n_pos : int
  action : int list -> int -> int -> unit
*)
let enumerate poset n_pos action =
  let rec aux = function
    |([],chosen,_,n_upper,n_lower,_,_) -> action chosen n_upper n_lower
    |(a::candidate, chosen, banned, n_upper, n_lower, c_upper, c_lower) ->
      (* a is the next element to decide to put in the set T^+ \cap T^- or not
        candidate is the list of elements waiting
        chosen is the list of elements chosen to be in the subset T^+ \cap T^-
        banned is the list of elements that we decided not to put into this intersection
        n_upper is the cardinality of the upper set generated by chosen
        n_lower is the cardinality of the lower set generated by chosen
        c_upper is the complement of the upper set generated by chosen in \Phi^+
        c_lower is the complement of the lower set generated by chosen in \dot\Phi
      *)
      begin
      (* If choosing a does not break the fact that chosen is the intersection of an upper set with a lower set, we may choose a *)
      if List.for_all (fun x -> List.for_all (fun y -> not (poset.(x).(y) && poset.(y).(a))) banned) chosen then
        begin
        let n_upper_incr, c_upper_rest = filter_count (function x -> poset.(a).(x)) c_upper
        and n_lower_incr, c_lower_rest = filter_count (function x -> poset.(x).(a)) c_lower in
        aux (candidate, a::chosen, banned, n_upper + n_upper_incr, n_lower + n_lower_incr, c_upper_rest, c_lower_rest)
        end;
      (* We may ban a *)
      aux (candidate, chosen, a::banned, n_upper, n_lower, c_upper, c_lower)
      end
  in
  let c_upper = List.init n_pos (function i -> i)
  and c_lower = List.init (Array.length poset) (function i -> i) in
  aux (c_upper, [], [], 0, 0, c_upper, c_lower)

let main t =
  let poset = Dynkin.matrix_ext_root_poset t
  and rank = Dynkin.rank t in
  let dim = Array.length poset in
  let n_pos_roots = (dim - rank) / 2 in
  Printf.printf "Dimension of the Lie algebra : %d\n" dim;
  Printf.printf "Number of positive roots : %d\n" n_pos_roots;
  flush stdout;
  Printf.printf "Enumerating all posibilities for T^+ \\cap T^-\n";
  flush stdout;
  let counter = ref 0 in
  let check subset n_upper n_lower =
    (*
    Printf.printf "{";
    List.iter (Printf.printf "%d, ") subset;
    Printf.printf "}. n_upper = %d, n_lower = %d\n" n_upper n_lower;
    *)
    incr counter;
    let n_cap = List.length subset in
    (* check inequality *)
    if dim * n_cap > n_upper * n_lower then raise (Counter_Example subset);
    (* check equality case *)
    if dim * n_cap = n_upper * n_lower && n_upper <> 0 && n_lower <> 0 && n_lower <> dim then raise (Counter_Example subset)
  in
  enumerate poset n_pos_roots check;
  Printf.printf "Done. Checked %d possibilities for T^+ \\cap T^-\n" !counter

let _ = main (Dynkin.of_string Sys.argv.(1))
